import os
import numpy as np
import matplotlib.pyplot as plt
import matplotlib

# 设置后端（应在 pyplot 导入后尽早设置）
matplotlib.use('TkAgg')

directory = r'./dataset-resized'

X = []  # 特征：每个样本是一张展平的灰度图 (height * width,)
y = []  # 标签：每个样本对应的类别名称

for class_name in os.listdir(directory):
    class_dir = os.path.join(directory, class_name)
    if os.path.isdir(class_dir):
        for img_name in os.listdir(class_dir):
            img_path = os.path.join(class_dir, img_name)
            try:
                im = plt.imread(img_path)

                # 处理不同图像格式（如 RGBA）
                if im.ndim == 3 and im.shape[2] in [3, 4]:
                    # 转为灰度图：加权平均
                    if im.shape[2] == 4:
                        # 如果有 alpha 通道，只使用 RGB 三通道
                        rgb = im[:, :, :3]
                    else:
                        rgb = im
                    gray_img = 0.299 * rgb[:, :, 0] + 0.587 * rgb[:, :, 1] + 0.114 * rgb[:, :, 2]
                elif im.ndim == 2:
                    gray_img = im  # 已经是灰度图
                else:
                    continue  # 不支持的维度

                # 展平图像为一维向量
                flat_img = gray_img.flatten()
                X.append(flat_img)
                y.append(class_name)

            except Exception as e:
                print(f"无法读取图像 {img_path}: {e}")

# 转换为 numpy 数组
X = np.array(X)  # shape: (n_samples, img_height * img_width)
y = np.array(y)  # shape: (n_samples,)

print("总样本数:", X.shape[0])
print("每张图像特征数（分辨率）:", X.shape[1])
print("类别数:", len(np.unique(y)))
print("类别:", np.unique(y))

# 保存
np.save('./X.npy', X)  # 大写 X 更合适
np.save('./y.npy', y)

print("数据保存完成：X.npy 和 y.npy")